import numpy as np
from util.TDataGenerator import DataGenerator

def load_data(filename):
    data = np.loadtxt(filename, delimiter=',')
    X = data[:, :-1]  # 特征
    y = data[:, -1]   # 目标值
    return X, y

def linear_regression(X, y):
    # 添加偏置项
    X_b = np.c_[np.ones((X.shape[0], 1)), X]  # 在特征矩阵中添加一列1
    # 计算线性方程的参数
    theta_best = np.linalg.inv(X_b.T.dot(X_b)).dot(X_b.T).dot(y)
    return theta_best

def compute_loss(X, y, theta):
    X_b = np.c_[np.ones((X.shape[0], 1)), X]  # 添加偏置项
    predictions = X_b.dot(theta)
    loss = np.mean((predictions - y) ** 2)  # 均方误差
    return loss

if __name__ == "__main__":
    num_rows = 5  # 指定行数
    num_cols = 3  # 指定列数

    generator = DataGenerator(num_rows, num_cols)
    filename = 'train_output/trainNum01.txt'  # 文件名

    # 写入数据到文件
    generator.write_data_to_file(filename)

    # 从文件读取数据
    X, y = load_data(filename)

    # 拟合线性方程
    theta = linear_regression(X, y)

    # 计算损失函数
    loss = compute_loss(X, y, theta)

    # 打印线性方程的参数和损失函数
    print("线性方程的参数:")
    print(theta)
    print("损失函数 (均方误差):")
    print(loss)